import importlib.util
from pathlib import Path

import numpy as np
import pandas as pd


def load_shiftshare_module(root: Path):
    path = root / "analysis" / "19_iv_shiftshare_rollout_absorb.py"
    spec = importlib.util.spec_from_file_location("shiftshare", path)
    mod = importlib.util.module_from_spec(spec)
    assert spec.loader is not None
    spec.loader.exec_module(mod)
    return mod


def to_latex_tabular(df: pd.DataFrame) -> str:
    cols = [
        "label",
        "nobs",
        "coef_lowEdu",
        "se_lowEdu",
        "coef_highEdu",
        "coef_deltaHighMinusLow",
        "se_deltaHighMinusLow",
        "fsF_netusoft",
        "fsF_netusoft_x_edu",
        "AR_pvalue",
    ]
    header = [
        "Omit country",
        "N",
        "b(LowEdu)",
        "se(LowEdu)",
        "b(HighEdu)",
        "Delta(High-Low)",
        "se(Delta)",
        "F(net)",
        "F(net$\\times$edu)",
        "AR p",
    ]
    d = df[cols].copy()
    d["nobs"] = d["nobs"].map(lambda x: f"{x:.0f}" if pd.notna(x) else "")
    for c in ["coef_lowEdu", "se_lowEdu", "coef_highEdu", "coef_deltaHighMinusLow", "se_deltaHighMinusLow"]:
        d[c] = d[c].map(lambda x: f"{x:.4f}" if pd.notna(x) else "")
    for c in ["fsF_netusoft", "fsF_netusoft_x_edu"]:
        d[c] = d[c].map(lambda x: f"{x:.2f}" if pd.notna(x) else "")
    d["AR_pvalue"] = d["AR_pvalue"].map(lambda x: f"{x:.4g}" if pd.notna(x) else "")

    out = []
    out.append("\\begin{tabular}{lrrrrrrrrr}")
    out.append("\\toprule")
    out.append(" & ".join(header) + " \\\\")
    out.append("\\midrule")
    for _, row in d.iterrows():
        out.append(" & ".join(str(row[c]) for c in cols) + " \\\\")
    out.append("\\bottomrule")
    out.append("\\end{tabular}")
    return "\n".join(out)


def main() -> None:
    root = Path(__file__).resolve().parents[1]
    out_dir = root / "outputs"
    out_dir.mkdir(parents=True, exist_ok=True)
    paper_tables = root / "paper_joc" / "tables"
    paper_tables.mkdir(parents=True, exist_ok=True)

    mod = load_shiftshare_module(root)
    df = mod.build_dataset(root)

    outcomes = [
        ("y_part_index", "Participation index"),
        ("y_vote", "Vote (yes/no)"),
    ]
    controls = ["agea", "gndr", "hinctnta"]
    countries = sorted(df["cntry"].dropna().unique().tolist())

    rows = []
    # Full-sample rows
    for y, label in outcomes:
        res = mod.fit_iv_interaction_absorb(df, y, controls=controls)
        r = mod.summarize_result(res, y=y, label=f"ALL ({label})")
        rows.append(r)

    for cntry in countries:
        d = df[df["cntry"] != cntry].copy()
        for y, label in outcomes:
            try:
                res = mod.fit_iv_interaction_absorb(d, y, controls=controls)
            except Exception:
                continue
            r = mod.summarize_result(res, y=y, label=f"{cntry} ({label})")
            rows.append(r)

    tab = pd.DataFrame(rows)
    (out_dir / "iv_shiftshare_leave_one_country_out.csv").write_text(tab.to_csv(index=False), encoding="utf-8")
    (out_dir / "iv_shiftshare_leave_one_country_out.md").write_text(mod.to_markdown_table(tab) + "\n", encoding="utf-8")
    (paper_tables / "iv_shiftshare_leave_one_country_out.tex").write_text(to_latex_tabular(tab) + "\n", encoding="utf-8")

    print("Wrote LOCO table to paper_joc/tables/iv_shiftshare_leave_one_country_out.tex")


if __name__ == "__main__":
    main()

